# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from experimental.prophet.graph import graph
from experimental.prophet.graph import graph_def_pb2
from experimental.prophet.ops.util import (to_tuple, to_data_type_proto, from_data_type_proto)


def concat(*vars, dtype=None, node_name=None, namespace=None):
  vars = [graph.as_variable(v) for v in to_tuple(vars)]

  # Determines an output dtype and casts all inputs.
  if dtype is None:
    dtype = graph_def_pb2.DATA_TYPE_BOOL
    for v in vars:
      var_dtype = to_data_type_proto(v.dtype)
      if dtype == graph_def_pb2.DATA_TYPE_BOOL:
        dtype = var_dtype
      elif dtype == graph_def_pb2.DATA_TYPE_INT64 and var_dtype == graph_def_pb2.DATA_TYPE_DOUBLE:
        dtype = var_dtype
  vars = [v.cast(dtype) for v in vars]

  node = graph.create_node('Concat',
                           vars,
                           name=node_name,
                           namespace=namespace,
                           attrs={'dtype': to_data_type_proto(dtype)})
  return node.outputs[0]
